// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/connectors/hive/storage_adapters/s3fs/RegisterS3FileSystem.h> // @manual

#ifdef POLLUX_ENABLE_S3
#include <pollux/common/base/StatsReporter.h>
#include <pollux/connectors/hive/storage_adapters/s3fs/S3Config.h> // @manual
#include <pollux/connectors/hive/storage_adapters/s3fs/S3Counters.h> // @manual
#include <pollux/connectors/hive/storage_adapters/s3fs/S3FileSystem.h> // @manual
#include <pollux/connectors/hive/storage_adapters/s3fs/S3Util.h> // @manual
#include <pollux/dwio/common/FileSink.h>
#endif

namespace kumo::pollux::filesystems {

#ifdef POLLUX_ENABLE_S3
using FileSystemMap = melon::Synchronized<
    std::unordered_map<std::string, std::shared_ptr<FileSystem>>>;

/// Multiple S3 filesystems are supported.
/// Key is the endpoint value specified in the config using hive.s3.endpoint.
/// If the endpoint is empty, it will default to AWS S3 Library.
/// Different S3 buckets can be accessed with different client configurations.
/// This allows for different endpoints, data read and write strategies.
/// The bucket specific option is set by replacing the hive.s3. prefix on an
/// option with hive.s3.bucket.BUCKETNAME., where BUCKETNAME is the name of the
/// bucket. When connecting to a bucket, all options explicitly set will
/// override the base hive.s3. values.

FileSystemMap& fileSystems() {
  static FileSystemMap instances;
  return instances;
}

CacheKeyFn cacheKeyFunc;

std::shared_ptr<FileSystem> fileSystemGenerator(
    std::shared_ptr<const config::ConfigBase> properties,
    std::string_view s3Path) {
  std::string cacheKey, bucketName, key;
  getBucketAndKeyFromPath(getPath(s3Path), bucketName, key);
  if (!cacheKeyFunc) {
    cacheKey = S3Config::cacheKey(bucketName, properties);
  } else {
    cacheKey = cacheKeyFunc(properties, s3Path);
  }

  // Check if an instance exists with a read lock (shared).
  auto fs = fileSystems().withRLock(
      [&](auto& instanceMap) -> std::shared_ptr<FileSystem> {
        auto iterator = instanceMap.find(cacheKey);
        if (iterator != instanceMap.end()) {
          return iterator->second;
        }
        return nullptr;
      });
  if (fs != nullptr) {
    return fs;
  }

  return fileSystems().withWLock(
      [&](auto& instanceMap) -> std::shared_ptr<FileSystem> {
        // Repeat the checks with a write lock.
        auto iterator = instanceMap.find(cacheKey);
        if (iterator != instanceMap.end()) {
          return iterator->second;
        }

        auto logLevel =
            properties->get(S3Config::kS3LogLevel, std::string("FATAL"));
        std::optional<std::string> logLocation =
            static_cast<std::optional<std::string>>(
                properties->get<std::string>(S3Config::kS3LogLocation));
        initializeS3(logLevel, logLocation);
        auto fs = std::make_shared<S3FileSystem>(bucketName, properties);
        instanceMap.insert({cacheKey, fs});
        return fs;
      });
}

std::unique_ptr<pollux::dwio::common::FileSink> s3WriteFileSinkGenerator(
    const std::string& fileURI,
    const pollux::dwio::common::FileSink::Options& options) {
  if (isS3File(fileURI)) {
    auto fileSystem =
        filesystems::getFileSystem(fileURI, options.connectorProperties);
    return std::make_unique<dwio::common::WriteFileSink>(
        fileSystem->openFileForWrite(fileURI, {{}, options.pool, std::nullopt}),
        fileURI,
        options.metricLogger,
        options.stats);
  }
  return nullptr;
}
#endif

void registerS3FileSystem(CacheKeyFn identityFunction) {
#ifdef POLLUX_ENABLE_S3
  fileSystems().withWLock([&](auto& instanceMap) {
    if (instanceMap.empty()) {
      cacheKeyFunc = identityFunction;
      registerFileSystem(isS3File, std::function(fileSystemGenerator));
      dwio::common::FileSink::registerFactory(
          std::function(s3WriteFileSinkGenerator));
    }
  });
#endif
}

void finalizeS3FileSystem() {
#ifdef POLLUX_ENABLE_S3
  bool singleUseCount = true;
  fileSystems().withWLock([&](auto& instanceMap) {
    for (const auto& [id, fs] : instanceMap) {
      singleUseCount &= (fs.use_count() == 1);
    }
    POLLUX_CHECK(singleUseCount, "Cannot finalize S3FileSystem while in use");
    instanceMap.clear();
  });

  finalizeS3();
#endif
}

void registerS3Metrics() {
#ifdef POLLUX_ENABLE_S3
  DEFINE_METRIC(kMetricS3ActiveConnections, pollux::StatType::SUM);
  DEFINE_METRIC(kMetricS3StartedUploads, pollux::StatType::COUNT);
  DEFINE_METRIC(kMetricS3FailedUploads, pollux::StatType::COUNT);
  DEFINE_METRIC(kMetricS3SuccessfulUploads, pollux::StatType::COUNT);
  DEFINE_METRIC(kMetricS3MetadataCalls, pollux::StatType::COUNT);
  DEFINE_METRIC(kMetricS3GetObjectCalls, pollux::StatType::COUNT);
  DEFINE_METRIC(kMetricS3GetObjectErrors, pollux::StatType::COUNT);
  DEFINE_METRIC(kMetricS3GetMetadataErrors, pollux::StatType::COUNT);
  DEFINE_METRIC(kMetricS3GetObjectRetries, pollux::StatType::COUNT);
  DEFINE_METRIC(kMetricS3GetMetadataRetries, pollux::StatType::COUNT);
#endif
}

void registerAWSCredentialsProvider(
    const std::string& providerName,
    const AWSCredentialsProviderFactory& provider) {
#ifdef POLLUX_ENABLE_S3
  registerCredentialsProvider(providerName, provider);
#endif
}

} // namespace kumo::pollux::filesystems
